import numpy as np
import pydicom
from pydicom import dcmread


def compute_translation_vector(params):
    """计算平移向量 T = [tx, ty, tz]"""
    theta_z = np.deg2rad(params["rao_lao_angle"])
    theta_x = np.deg2rad(params["cran_caud_angle"])
    sad = params["sad"]
    
    # 初始位置：未旋转时的源位置 (0, 0, SAD)
    T0 = np.array([0, 0, sad])
    
    # 绕Z轴旋转矩阵（RAO/LAO）
    Rz = np.array([
        [np.cos(theta_z), -np.sin(theta_z), 0],
        [np.sin(theta_z), np.cos(theta_z), 0],
        [0, 0, 1]
    ])
    
    # 绕X轴旋转矩阵（CRAN/CAUD）
    Rx = np.array([
        [1, 0, 0],
        [0, np.cos(theta_x), -np.sin(theta_x)],
        [0, np.sin(theta_x), np.cos(theta_x)]
    ])
    
    # 复合旋转并计算平移向量
    T = Rz @ Rx @ T0
    return T.reshape(3, 1)  # 返回3x1向量

def load_dicom_parameters(dicom_path):
    """从DICOM文件中提取关键参数"""
    ds = dcmread(dicom_path)
    
    # 提取内参相关参数
    sid = ds.get((0x0018, 0x1110), 1000.0).value  # 默认假设SID=1000 mm
    pixel_spacing = [0.214488,0.214488]#ds.get((0x0018, 0x1164), [0.2, 0.2]).value  # 默认假设0.2 mm/pixel
    sad = ds.get((0x0018, 0x1111), 1000.0).value
    rows = ds.Rows
    cols = ds.Columns
    
    # 提取外参相关参数（部分标签可能不存在或为私有）
    try:
        rao_lao_angle = ds[0x0018, 0x1510].value  # RAO/LAO角度（度）
        cran_caud_angle = ds[0x0018, 0x1511].value  # CRAN/CAUD角度（度）
    except KeyError:
        rao_lao_angle = 0.0
        cran_caud_angle = 0.0
    
    return {
        "image":np.array(ds.pixel_array),
        "sad":sad,
        "sid": sid,
        "pixel_spacing": pixel_spacing,
        "image_size": (rows, cols),
        "rao_lao_angle": rao_lao_angle,
        "cran_caud_angle": cran_caud_angle
    }

def compute_intrinsic_matrix(params):
    """计算内参矩阵 K"""
    fx = params["sid"] / params["pixel_spacing"][0]  # 焦距fx = SID / 像素宽
    fy = params["sid"] / params["pixel_spacing"][1]  # 焦距fy = SID / 像素高
    u0 = params["image_size"][1] / 2.0  # 主点u0（列中心）
    v0 = params["image_size"][0] / 2.0  # 主点v0（行中心）
    
    K = np.array([
        [fx, 0, u0],
        [0, fy, v0],
        [0, 0, 1]
    ])
    return K

def compute_extrinsic_matrix(params):
    """计算外参矩阵 [R|T]（假设平移向量T为0，需根据实际情况调整）"""
    # 将角度转换为弧度
    theta_z = np.deg2rad(params["rao_lao_angle"])  # RAO/LAO绕Z轴旋转
    theta_x = np.deg2rad(params["cran_caud_angle"])  # CRAN/CAUD绕X轴旋转
    
    # 计算旋转矩阵（假设旋转顺序为Z-X）
    Rz = np.array([
        [np.cos(theta_z), -np.sin(theta_z), 0],
        [np.sin(theta_z), np.cos(theta_z), 0],
        [0, 0, 1]
    ])
    
    Rx = np.array([
        [1, 0, 0],
        [0, np.cos(theta_x), -np.sin(theta_x)],
        [0, np.sin(theta_x), np.cos(theta_x)]
    ])
    
    R = Rz @ Rx  # 复合旋转矩阵
    
    # 假设平移向量T为0（需根据SAD和几何关系调整）
    T = compute_translation_vector(params)
    
    # return np.hstack([R, T])
    return R,T
#%%
if __name__ == "__main__":
    pass
    #%%
    import matplotlib.pyplot as plt
    dicom_path = "../resource/IM000000"  # 替换为你的DICOM文件路径
    params = load_dicom_parameters(dicom_path)
    plt.imshow(params['image'][22,:,:],cmap='gray')
    print(params)
    # 计算内参矩阵
    K = compute_intrinsic_matrix(params)
    print("内参矩阵 K:\n", K)
    
    # 计算外参矩阵
    R_T = compute_extrinsic_matrix(params)
    print("\n外参矩阵 [R|T]:\n", R_T)
# %%
